#include <iostream>
#include <chrono>
#include <thread>
#include <csignal>
#include <unordered_map>

#include "utils/logger.h"
#include "utils/command_line_parser.h"
#include "core/device_manager.h"
#include "core/packet_dispatcher.h"
#include "filters/types.h"
#include "filters/simple_filter_engine.h"
#include "filters/filter_worker.h"

std::atomic<bool> g_running{true};
void signalHandler(int signal) {
    g_running = false;
}


int main(int argc, char* argv[]) {
    /// Register semaphores
    std::signal(SIGINT, signalHandler);    /// Ctrl + c
    std::signal(SIGTERM, signalHandler);   /// kill

    std::cout << "Hello hp filter!" << std::endl;
    CommandLineParser parser;
    if (!parser.parse(argc, argv)) {
        return 1;
    }
    if (parser.hasHelp()) {
        parser.printHelp();
        return 0;
    }
    if (parser.hasVersion()) {
        parser.printVersion();
        return 0;
    }
    std::string log_level = parser.getLogLevel();
    std::string log_file = parser.getLogFile();
    if (parser.isConsoleOnly()) {
        log_file = "";
    }
    if (log_level == "debug" || log_level == "Debug" || log_level == "DEBUG") {
        Logger::getInstance().init(LogLevel::DEBUG, log_file);
    } else if (log_level == "info" || log_level == "Info" || log_level == "INFO") {
        Logger::getInstance().init(LogLevel::INFO, log_file);
    } else if (log_level == "warning" || log_level == "Warning" || log_level == "WARNING") {
        Logger::getInstance().init(LogLevel::WARNING, log_file);
    } else if (log_level == "error" || log_level == "Error" || log_level == "ERROR") {
        Logger::getInstance().init(LogLevel::ERROR, log_file);
    } else {
        std::cerr << "Failed to set log level, use \"info\" by default." << std::endl;
        Logger::getInstance().init(LogLevel::INFO, log_file);
    }
    auto& manager = DeviceManager::getInstance();
    if (!manager.initialize()) {
        LOG_ERROR("Failed to initialize device manager.");
    }
    if (parser.listDevices()) {
        std::vector<std::string> list = manager.listAvailableDevices();
        for (const auto& deviceName : list) {
            std::cout << deviceName << std::endl;
        }
        return 0;
    }

    /// Default configuration
    int workers = 4;
    int bufferSize = 8 * 1024 * 1024;
    int batchSize = 4;
    int timeout = 1000;

    if (parser.getWorkers() != 0) {
        workers = parser.getWorkers();
    }
    if (parser.getBufferSize() != 0) {
        bufferSize = parser.getBufferSize();
    }
    if (parser.getBatchSize() != 0) {
        batchSize = parser.getBatchSize();
    }
    if (parser.getTimeOut() != 0) {
        timeout = parser.getTimeOut();
    }
    std::vector<std::string> devices = parser.getDevices();
    std::unordered_map<std::string, std::shared_ptr<PacketDispatcher>> dispatchers;
    std::vector<std::unique_ptr<FilterThreadPool>> filter_pools;
    for (auto& device : devices) {
        auto filterEngine = std::make_shared<SimpleFilterEngine>();
        FilterRule rule1;
        rule1.id = 1;
        rule1.name = "Allow Localhost TCP";
        rule1.condition.protocol = 6;
        rule1.condition.src_ip = IPRange::fromCIDR("127.0.0.1/32").value();
        rule1.action = FilterAction::ACCEPT;
        rule1.priority = 1;
        rule1.enable = true;
        filterEngine->addRule(rule1);
        FilterRule rule2;
        rule2.id = 2;
        rule2.name = "Block ICMP";
        rule2.condition.protocol = 1;
        rule2.action = FilterAction::DROP;
        rule2.priority = 10;
        rule2.enable = true;
        filterEngine->addRule(rule2);
        auto dispatcher = std::make_shared<PacketDispatcher>(workers, DispatchStrategy::ROUND_ROBIN);
        dispatchers[device] = dispatcher;
        filter_pools.emplace_back(std::make_unique<FilterThreadPool>(workers, filterEngine, dispatcher));
        auto d = manager.createDevice(device);
        d->setPromiscuousMode(parser.isPromiscuous());
        d->setBufferSize(bufferSize);
        if (!parser.getFilter().empty()) {
            d->setBPFFilter(parser.getFilter());
        }
        d->setTimeout(timeout);
        d->setPacketCallback([dispatchers, device](Packet&& packet){
            LOG_DEBUG("One packet arrived.");
            dispatchers.at(device)->dispatch(std::move(packet));
        });
        filter_pools.back()->start();
        d->startCapture();
    }
    /// Wait exit signal.
    while (g_running) {
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
    for (auto& device : devices) {
        manager.removeDevice(device);
    }
    for (auto& pool : filter_pools){
        pool->stop();
    }
    return 0;
}